import sys
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtWidgets import (QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, 
                           QPushButton, QHBoxLayout, QMessageBox, QFileDialog, QProgressBar,
                           QLineEdit, QDialog)
from PyQt5.QtGui import QImage, QPixmap, QIcon
import cv2
from ultralytics import YOLO
from openai import OpenAI
import json

# 添加登录窗口类
class LoginDialog(QDialog):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("登录")
        self.setFixedSize(300, 200)
        
        layout = QVBoxLayout()
        
        # 用户名输入
        username_layout = QHBoxLayout()
        username_label = QLabel("用户名:")
        self.username_input = QLineEdit()
        username_layout.addWidget(username_label)
        username_layout.addWidget(self.username_input)
        
        # 密码输入
        password_layout = QHBoxLayout()
        password_label = QLabel("密码:")
        self.password_input = QLineEdit()
        self.password_input.setEchoMode(QLineEdit.Password)  # 设置密码隐藏
        password_layout.addWidget(password_label)
        password_layout.addWidget(self.password_input)
        
        # 登录按钮
        self.login_button = QPushButton("登录")
        self.login_button.clicked.connect(self.check_login)
        
        # 添加所有组件到主布局
        layout.addLayout(username_layout)
        layout.addLayout(password_layout)
        layout.addWidget(self.login_button)
        
        self.setLayout(layout)
        
        # 设置默认的用户名和密码
        self.correct_username = "admin"
        self.correct_password = "123456"
        
    def check_login(self):
        username = self.username_input.text()
        password = self.password_input.text()
        
        if username == self.correct_username and password == self.correct_password:
            self.accept()  # 登录成功
        else:
            QMessageBox.warning(self, "错误", "用户名或密码错误！")

# 定义一个工作类，用于加载模型和进行检测
class Worker:
    def __init__(self):
        self.model = None  # 初始化模型为 None

    def load_model(self):
        # 打开文件对话框选择模型文件
        model_path, _ = QFileDialog.getOpenFileName(None, "选择模型文件", "", "模型文件 (*.pt)")
        if model_path:
            self.model = YOLO(model_path)  # 加载 YOLO 模型
            return self.model is not None  # 返回模型是否成功加载
        return False

    def detect_image(self, image):
        # 使用模型进行图像检测
        results = self.model.predict(image)
        return results

# 定义主窗口类
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("李健驹的农作物病害检测系统")  # 设置窗口标题
        self.setGeometry(300, 150, 800, 400)  # 设置窗口大小和位置

        # 创建两个 QLabel 分别显示左右图像
        self.label1 = QLabel()
        self.label1.setAlignment(Qt.AlignCenter)  # 设置对齐方式为居中
        self.label1.setMinimumSize(580, 450)  # 设置最小尺寸
        self.label1.setStyleSheet('border:3px solid #6950a1; background-color: black;')  # 设置样式

        self.label2 = QLabel()
        self.label2.setAlignment(Qt.AlignCenter)
        self.label2.setMinimumSize(580, 450)
        self.label2.setStyleSheet('border:3px solid #6950a1; background-color: black;')

        # 创建布局
        layout = QVBoxLayout()
        hbox_video = QHBoxLayout()
        hbox_video.addWidget(self.label1)  # 左侧显示原始图像
        hbox_video.addWidget(self.label2)  # 右侧显示检测后的图像
        layout.addLayout(hbox_video)
        self.worker = Worker()  # 创建 Worker 实例

        # 创建按钮布局
        hbox_buttons = QHBoxLayout()

        # 添加模型选择按钮
        self.load_model_button = QPushButton("📁模型选择")
        self.load_model_button.clicked.connect(self.load_model)  # 连接按钮点击事件
        self.load_model_button.setFixedSize(120, 30)  # 设置按钮大小
        hbox_buttons.addWidget(self.load_model_button)

        # 添加图片检测按钮
        self.image_detect_button = QPushButton("💾图片检测")
        self.image_detect_button.clicked.connect(self.detect_image)
        self.image_detect_button.setEnabled(False)  # 初始状态禁用
        self.image_detect_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.image_detect_button)

        # 添加显示检测物体按钮
        self.display_objects_button = QPushButton("🔍显示检测物体")
        self.display_objects_button.clicked.connect(self.show_detected_objects)
        self.display_objects_button.setEnabled(False)
        self.display_objects_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.display_objects_button)

        # 添加视频流检测按钮
        self.video_detect_button = QPushButton("📹视频流检测")
        self.video_detect_button.clicked.connect(self.toggle_video_detection)
        self.video_detect_button.setEnabled(False)
        self.video_detect_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.video_detect_button)

        # 添加暂停按钮
        self.pause_button = QPushButton("⏸暂停")
        self.pause_button.clicked.connect(self.pause_video)
        self.pause_button.setEnabled(False)
        self.pause_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.pause_button)

        # 添加继续播放按钮
        self.resume_button = QPushButton("▶️继续")
        self.resume_button.clicked.connect(self.resume_video)
        self.resume_button.setEnabled(False)
        self.resume_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.resume_button)

        # 添加导出按钮
        self.export_button = QPushButton("💾导出结果")
        self.export_button.clicked.connect(self.export_results)
        self.export_button.setEnabled(False)
        self.export_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.export_button)

        # 添加 AI 分析按钮
        self.ai_analysis_button = QPushButton("🤖 AI分析")
        self.ai_analysis_button.clicked.connect(self.perform_ai_analysis)
        self.ai_analysis_button.setEnabled(False)
        self.ai_analysis_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.ai_analysis_button)

        # 添加进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setFixedSize(580, 20)
        layout.addWidget(self.progress_bar)

        # 添加退出按钮
        self.exit_button = QPushButton("❌退出")
        self.exit_button.clicked.connect(self.exit_application)
        self.exit_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.exit_button)

        layout.addLayout(hbox_buttons)
        central_widget = QWidget()
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        self.current_results = None  # 当前检测结果

        self.timer = QTimer()  # 定时器用于视频帧处理
        self.timer.timeout.connect(self.process_video_frame)
        self.cap = None  # 视频捕获对象

        # 初始化 DeepSeek AI 客户端
        self.client = OpenAI(
            base_url="https://api.deepseek.com/v1",
            api_key="sk-8fedb997b12e4c0c85012a0abce85002"  # 替换为你的 API key
        )

    def detect_image(self):
        # 打开文件对话框选择图片文件
        image_path, _ = QFileDialog.getOpenFileName(None, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png)")
        if image_path:
            image = cv2.imread(image_path)  # 读取图像
            if image is not None:
                self.current_results = self.worker.detect_image(image)  # 检测图像
                if self.current_results:
                    annotated_image = self.current_results[0].plot()  # 绘制检测结果
                    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转换为 RGB
                    height1, width1, channel1 = image_rgb.shape
                    bytesPerLine1 = 3 * width1
                    qimage1 = QImage(image_rgb.data, width1, height1, bytesPerLine1, QImage.Format_RGB888)
                    pixmap1 = QPixmap.fromImage(qimage1)
                    self.label1.setPixmap(pixmap1.scaled(self.label1.size(), Qt.KeepAspectRatio))

                    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)  # 转换为 RGB
                    height2, width2, channel2 = annotated_image.shape
                    bytesPerLine2 = 3 * width2
                    qimage2 = QImage(annotated_image.data, width2, height2, bytesPerLine2, QImage.Format_RGB888)
                    pixmap2 = QPixmap.fromImage(qimage2)
                    self.label2.setPixmap(pixmap2.scaled(self.label2.size(), Qt.KeepAspectRatio))

                    self.export_button.setEnabled(True)  # 启用导出按钮
                    self.ai_analysis_button.setEnabled(True)  # 启用 AI 分析按钮

    def show_detected_objects(self):
        if self.current_results:
            det_info = self.current_results[0].boxes.cls  # 获取检测到的物体类别
            object_count = len(det_info)  # 计算物体数量
            object_info = f"识别到的物体总个数：{object_count}\n"
            object_dict = {}
            class_names_dict = self.current_results[0].names  # 获取类别名称
            for class_id in det_info:
                class_name = class_names_dict[int(class_id)]
                if class_name in object_dict:
                    object_dict[class_name] += 1
                else:
                    object_dict[class_name] = 1
            sorted_objects = sorted(object_dict.items(), key=lambda x: x[1], reverse=True)
            for obj_name, obj_count in sorted_objects:
                object_info += f"{obj_name}: {obj_count}\n"
            self.show_message_box("识别结果", object_info)  # 显示检测结果
        else:
            self.show_message_box("识别结果", "未检测到物体")

    def show_message_box(self, title, message):
        msg_box = QMessageBox(self)  # 创建消息框
        msg_box.setWindowTitle(title)  # 设置标题
        msg_box.setText(message)  # 设置消息内容
        msg_box.exec_()  # 显示消息框

    def load_model(self):
        if self.worker.load_model():
            self.image_detect_button.setEnabled(True)  # 启用图片检测按钮
            self.video_detect_button.setEnabled(True)  # 启用视频检测按钮
            self.display_objects_button.setEnabled(True)  # 启用显示检测物体按钮

    def exit_application(self):
        sys.exit()  # 终止程序运行

    def toggle_video_detection(self):
        if self.timer.isActive():
            self.timer.stop()  # 停止定时器
            if self.cap:
                self.cap.release()  # 释放视频捕获对象
            self.video_detect_button.setText("📹视频流检测")
            self.pause_button.setEnabled(False)  # 禁用暂停按钮
            self.resume_button.setEnabled(False)  # 禁用继续按钮
        else:
            video_path, _ = QFileDialog.getOpenFileName(None, "选择视频文件", "", "视频文件 (*.mp4 *.avi)")
            if video_path:
                self.cap = cv2.VideoCapture(video_path)  # 打开视频文件
                if self.cap.isOpened():
                    self.timer.start(30)  # 每30毫秒处理一帧
                    self.video_detect_button.setText("⏹停止检测")
                    self.pause_button.setEnabled(True)  # 启用暂停按钮
                    self.resume_button.setEnabled(False)  # 禁用继续按钮
                    self.progress_bar.setMaximum(int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)))  # 设置进度条最大值

    def process_video_frame(self):
        ret, frame = self.cap.read()  # 读取视频帧
        if ret:
            self.current_results = self.worker.detect_image(frame)  # 检测视频帧
            if self.current_results:
                annotated_frame = self.current_results[0].plot()  # 绘制检测结果
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # 转换为 RGB
                height1, width1, channel1 = frame_rgb.shape
                bytesPerLine1 = 3 * width1
                qimage1 = QImage(frame_rgb.data, width1, height1, bytesPerLine1, QImage.Format_RGB888)
                pixmap1 = QPixmap.fromImage(qimage1)
                self.label1.setPixmap(pixmap1.scaled(self.label1.size(), Qt.KeepAspectRatio))

                annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)  # 转换为 RGB
                height2, width2, channel2 = annotated_frame.shape
                bytesPerLine2 = 3 * width2
                qimage2 = QImage(annotated_frame.data, width2, height2, bytesPerLine2, QImage.Format_RGB888)
                pixmap2 = QPixmap.fromImage(qimage2)
                self.label2.setPixmap(pixmap2.scaled(self.label2.size(), Qt.KeepAspectRatio))

            # 更新进度条
            current_frame = int(self.cap.get(cv2.CAP_PROP_POS_FRAMES))
            self.progress_bar.setValue(current_frame)

            self.export_button.setEnabled(True)  # 启用导出按钮
            self.ai_analysis_button.setEnabled(True)  # 启用 AI 分析按钮
        else:
            self.timer.stop()  # 停止定时器
            self.cap.release()  # 释放视频捕获对象
            self.video_detect_button.setText("📹视频流检测")
            self.pause_button.setEnabled(False)  # 禁用暂停按钮
            self.resume_button.setEnabled(False)  # 禁用继续按钮

    def pause_video(self):
        if self.timer.isActive():
            self.timer.stop()  # 停止定时器
            self.pause_button.setEnabled(False)  # 禁用暂停按钮
            self.resume_button.setEnabled(True)  # 启用继续按钮

    def resume_video(self):
        if not self.timer.isActive() and self.cap:
            self.timer.start(30)  # 启动定时器
            self.pause_button.setEnabled(True)  # 启用暂停按钮
            self.resume_button.setEnabled(False)  # 禁用继续按钮

    def export_results(self):
        if self.current_results:
            options = QFileDialog.Options()
            file_path, _ = QFileDialog.getSaveFileName(self, "保存检测结果", "", "图片文件 (*.png);;视频文件 (*.mp4)", options=options)
            if file_path:
                if file_path.endswith('.png'):
                    # 导出图片
                    annotated_image = self.current_results[0].plot()
                    cv2.imwrite(file_path, annotated_image)  # 保存图片
                elif file_path.endswith('.mp4'):
                    # 导出视频
                    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                    out = cv2.VideoWriter(file_path, fourcc, 30.0, (self.label2.width(), self.label2.height()))
                    while self.cap.isOpened():
                        ret, frame = self.cap.read()
                        if not ret:
                            break
                        results = self.worker.detect_image(frame)
                        if results:
                            annotated_frame = results[0].plot()
                            out.write(annotated_frame)  # 写入视频帧
                    out.release()  # 释放视频写入对象
                QMessageBox.information(self, "导出成功", "检测结果已成功导出！")  # 显示导出成功消息

    # 添加 AI 分析方法
    def perform_ai_analysis(self):
        if not self.current_results:
            QMessageBox.warning(self, "警告", "请先进行检测！")
            return

        try:
            # 获取检测结果
            det_info = self.current_results[0].boxes.cls
            class_names_dict = self.current_results[0].names
            
            # 构建检测结果描述
            detected_objects = []
            object_dict = {}
            for class_id in det_info:
                class_name = class_names_dict[int(class_id)]
                if class_name in object_dict:
                    object_dict[class_name] += 1
                else:
                    object_dict[class_name] = 1
            
            for obj_name, count in object_dict.items():
                detected_objects.append(f"{obj_name}: {count}个")

            # 准备发送给 AI 的提示
            prompt = f"""
            我检测到了以下农作物病害情况：
            {', '.join(detected_objects)}
            
            请分析这些检测结果，并提供：
            1. 这些病害的严重程度评估
            2. 可能造成的影响
            3. 建议采取的防治措施
            4. 预防建议
            
            请用中文回答，并尽量详细专业。
            """

            # 调用 DeepSeek AI API
            response = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7,
                max_tokens=2000
            )

            # 显示 AI 分析结果
            analysis_result = response.choices[0].message.content
            self.show_ai_analysis_result(analysis_result)

        except Exception as e:
            QMessageBox.critical(self, "错误", f"AI分析过程中发生错误：{str(e)}")

    def show_ai_analysis_result(self, analysis):
        # 创建一个更大的消息框来显示分析结果
        msg_box = QMessageBox(self)
        msg_box.setWindowTitle("AI 分析报告")
        msg_box.setText(analysis)
        msg_box.setStyleSheet("QLabel{min-width: 500px;}")
        msg_box.exec_()

# 修改主程序入口
if __name__ == '__main__':
    app = QApplication(sys.argv)
    
    # 首先显示登录窗口
    login = LoginDialog()
    if login.exec_() == QDialog.Accepted:  # 如果登录成功
        window = MainWindow()  # 创建主窗口
        window.show()  # 显示主窗口
        sys.exit(app.exec_())
    else:
        sys.exit()  # 如果登录失败或取消登录，直接退出程序